<?php
namespace Mcp\Server\Transport;
use Mcp\Server\Transport\TransportInterface;
use Symfony\Component\Uid\Uuid;
use Psr\Log\LoggerAwareInterface;
use Psr\Log\LoggerAwareTrait;
use Swoole\Coroutine;
use Swoole\Coroutine\Http\Client;
use Swoole\WebSocket\Frame;

/**
 * Swoole WebSocket Client Transport (MCP Compliant).
 *
 * This transport connects to a remote WebSocket server (like the Xiaozhi MCP endpoint)
 * and behaves exactly like the StdioTransport, handling a single session.
 */
class SwooleWebSocketClientTransport implements TransportInterface, LoggerAwareInterface
{
    use LoggerAwareTrait;

    private string $url;
    private $client = null;
    private bool $isConnected = false;
    private ?Uuid $sessionId = null;
    private $messageListener = null;
    private $pingTimerId = null;
    private const PING_INTERVAL = 25;
    public function __construct(string $url)
    {
        $this->url = $url;
    }

    public function initialize(): void
    {
        $this->logInfo('Swoole WebSocket client transport initialized.', [
            'target_url' => $this->url,
        ]);
    }

    public function onMessage(callable $listener): void
    {
        $this->messageListener = $listener;
	    $this->logInfo('Message listener registered');
    }

    public function onSessionEnd(callable $listener): void
    {
        $this->sessionEndListener = $listener;
    }

    /**
     * 连接到远程 WebSocket 服务器并进入监听循环。
     * 此方法会阻塞当前进程（协程）。
     */
    public function listen(): mixed
    {
        $this->logInfo('listen');
        if ($this->isConnected) {
            $this->logInfo('Transport is already connected.');
            return null;
        }
        try {
            $this->connectToServer();
            $this->startPingTimer();
        } catch (\Throwable $e) {
            $this->logInfo('Failed to start WebSocket transport: ' . $e->getMessage(), [
                'exception' => $e
            ]);
            $this->handleClose();
        }
        return null;
    }

    /**
     * 发送数据到已连接的 WebSocket 服务器。
     * 严格遵循 SDK 预期，忽略 $context 参数。
     */
    public function send(string $data, array $context = []): void
    {
        // 关键改动：忽略 $context，只检查连接状态
        if (!$this->isConnected || !$this->client) {
            $this->logInfo('Cannot send data: WebSocket connection is not active.');
            return;
        }

        // 直接通过已建立的连接发送数据
        $this->client->push($data);
        $this->logInfo('Sent data via WebSocket client.'.$data);
    }

    public function close(): void
    {
        $this->handleClose();
    }

    public function getSessionId(): ?Uuid
    {
        return $this->sessionId;
    }

    /**
     * 实际执行连接操作。
     */
    private function connectToServer(): void
    {
        $this->logInfo('connectToServer');
        $parsedUrl = parse_url($this->url);
        if ($parsedUrl === false) {
            $this->logInfo('Invalid MCP WebSocket URL: {$this->url}');
            throw new \InvalidArgumentException("Invalid MCP WebSocket URL: {$this->url}");
        }
        $host = $parsedUrl['host'] ?? '127.0.0.1';
        $port = $parsedUrl['port'] ?? ($parsedUrl['scheme'] === 'wss' ? 443 : 80);
        $path = $parsedUrl['path'] ?? '/';
        if (!empty($parsedUrl['query'])) {
            $path .= '?' . $parsedUrl['query'];
        }
        $ssl = strtolower($parsedUrl['scheme'] ?? '') === 'wss';

        $this->client = new Client($host, $port, $ssl);
        $this->client->set([
            'timeout' => 10,
            'connect_timeout' => 5,
            'headers' => [
                'User-Agent' => 'MCP-PHP-SDK-Swoole-Client',
                'Upgrade' => 'websocket',
                'Connection' => 'Upgrade',
                'Sec-WebSocket-Key' => base64_encode(random_bytes(16)),
                'Sec-WebSocket-Version' => '13',
            ],
        ]);

        // Perform WebSocket handshake
        if (!$this->client->upgrade($path)) {
            $errMsg = swoole_strerror($this->client->errCode);
            $this->logError("Failed to upgrade connection to WebSocket.", [
                'host' => $host,
                'port' => $port,
                'path' => $path,
                'ssl' => $ssl,
                'error_code' => $this->client->errCode,
                'error_message' => $errMsg
            ]);
            throw new \RuntimeException("WebSocket upgrade failed: {$errMsg} (Code: {$this->client->errCode})");
        }

        $this->isConnected = true;
        $this->logInfo('Successfully upgraded to WebSocket protocol.');

        // Start a new coroutine to poll for messages, so it doesn't block the listen() method.
        Coroutine::create(function () {
            $this->startMessagePolling();
        });
    }

    /**
     * Manually polls for incoming WebSocket messages using recv().
     * This runs in a separate coroutine.
     */
    private function startMessagePolling(): void
    {
        $this->logInfo('Message polling coroutine started.');
        while ($this->isConnected && $this->client) {
            $frame = $this->client->recv(0.5); 
            if ($frame === false) {
                $errCode = $this->client->errCode;
                if ($errCode === SOCKET_ETIMEDOUT) {
                    continue;
                }
                $errMsg = swoole_strerror($errCode);
                $this->logError('Error receiving data from WebSocket server.', [
                    'error_code' => $errCode,
                    'error_message' => $errMsg
                ]);
                break;
            }

            if ($frame === '') {
                $this->logInfo('WebSocket connection closed by the server.');
                break;
            }
            if ($frame instanceof Frame && ($frame->opcode == SWOOLE_WEBSOCKET_OPCODE_TEXT || $frame->opcode == SWOOLE_WEBSOCKET_OPCODE_BINARY)) {
                $this->logInfo('Received message from WebSocket server.'.$frame->data);
                if ($this->messageListener) {
                    try {
                        call_user_func($this->messageListener, $frame->data, $this->sessionId);
                    } catch (\Throwable $e) {
                        $this->logError('An error occurred while invoking the message listener.'.$e);
                    }
                }
            } elseif ($frame instanceof Frame && $frame->opcode == SWOOLE_WEBSOCKET_OPCODE_PING) {
                $this->client->push('', SWOOLE_WEBSOCKET_OPCODE_PONG);
                $this->logInfo('Responded to server PING with PONG.');
            }
        }

        // If we exit the loop, the connection is no longer active
        $this->handleClose();
    }
     private function startPingTimer(): void
    {
        if ($this->pingTimerId !== null) {
            return;
        }
        $this->pingTimerId = \Swoole\Timer::tick(self::PING_INTERVAL * 1000, function () {
            if ($this->isConnected && $this->client) {
                try {
                    // Swoole WebSocket 客户端会自动处理 ping/pong
                    // 我们只需尝试发送一个空消息或利用 Swoole 的心跳机制
                    // 这里我们发送一个空的文本帧，服务器会自动回复 pong
                    $this->client->push(''); 
                    $this->logInfo("Sent ping to keep connection alive.");
                } catch (\Throwable $e) {
                    $this->logWarning("Failed to send ping: " . $e->getMessage());
                    // 发送失败，说明连接可能已断开，尝试关闭
                    $this->handleClose();
                }
            } else {
                // 如果连接已断开，清除计时器
                $this->stopPingTimer();
            }
        });
    }
    private function stopPingTimer(): void
    {
        if ($this->pingTimerId !== null) {
            \Swoole\Timer::clear($this->pingTimerId);
            $this->pingTimerId = null;
            $this->logInfo("Ping timer stopped.");
        }
    }
    /**
     * 处理连接关闭的统一逻辑。
     */
    private function handleClose(): void
    {
        if (!$this->isConnected) {
            return;
        }

        $this->isConnected = false;
        $sessionId = $this->sessionId ? $this->sessionId->toRfc4122(): '';
        $this->logInfo('Closing WebSocket connection.', ['session_id' => $sessionId]);
        $this->stopPingTimer();
        if ($this->client) {
            $this->client->close();
            $this->client = null;
        }
        $this->sessionId = null;
	    $this->logInfo('WebSocket connection closed and resources released.'.$sessionId);
    }
    /**
     * 外部设置会话ID（仅在 initialize 响应时调用）
     */
    public function setSessionId(Uuid $sessionId): void
    {
        $this->sessionId = $sessionId;
        $this->logInfo('Session ID set successfully.'.$sessionId->toRfc4122());
    }
    /**
     * 日志记录
     */
    public function logInfo(string $message): void
    {
        $this->log("[INFO] {$message}");
    }

    private function logWarning(string $message): void {
        $this->log("[WARNING] {$message}");
    }

    private function logError(string $message): void {
        $this->log("[ERROR] {$message}");
    }

    private function log(string $message): void
    {
        $logFile = ROOT_D . '/addons/jy_agentmcp/mcp_' . date('Ymd') . '.log';
        $dir = dirname($logFile);
        if (!is_dir($dir)) {
            mkdir($dir, 0755, true);
        }
        $content = date('Y-m-d H:i:s') . " {$message}\n";
        file_put_contents($logFile, $content, FILE_APPEND);
    }
}